package algorithm_demo.demo04.graph;

import java.util.*;

/**
 * 图：最小生成树算法之Kruskal
 * 1. 总是从权值最小的边开始考虑，依次考察权值依次变大的边
 * 2. 当前的边要么进入最小生成树的集合，要么丢弃
 * 3. 如果当前的边进入最小生成树的集合中不会形成环，就要当前边
 * 4. 如果当前的边进入最小生成树的集合中会形成环，就不要当前边
 * 5. 考察完所有边之后，最小生成树的集合也得到了。
 *
 * @author Api
 * @date 2023/2/19 10:23
 */
public class Code04_Kruskal {
    //并查集
    public static class UnionFind {
        public Map<Node, Node> parents;
        public Map<Node, Integer> size;

        public UnionFind() {
            parents = new HashMap<>();
            size = new HashMap<>();
        }

        public void makeSets(Collection<Node> nodes) {
            parents.clear();
            size.clear();
            for (Node node : nodes) {
                parents.put(node, node);
                size.put(node, 1);
            }
        }

        public Node findFather(Node cur) {
            Stack<Node> path = new Stack<>();
            while (cur != parents.get(cur)) {
                path.push(cur);
                cur = parents.get(cur);
            }
            while (!path.isEmpty()) {
                parents.put(path.pop(), cur);
            }
            return cur;
        }

        public boolean isSameSet(Node a, Node b) {
            return findFather(a) == findFather(b);
        }


        public void union(Node a, Node b) {
            if (a == null || b == null) {
                return;
            }
            Node aHead = findFather(a);
            Node bHead = findFather(b);
            if (aHead != bHead) {
                int aSetSize = size.get(aHead);
                int bSetSize = size.get(bHead);
                Node big = aSetSize >= bSetSize ? aHead : bHead;
                Node small = big == aHead ? bHead : aHead;
                parents.put(small, big);
                size.put(big, aSetSize + bSetSize);
                size.remove(small);
            }
        }
    }

    public static Set<Edge> kruskalMST(Graph graph) {
        UnionFind u = new UnionFind();
        u.makeSets(graph.nodes.values());
        Set<Edge> ans = new HashSet<>();
        PriorityQueue<Edge> heap = new PriorityQueue<>(Comparator.comparingInt(o -> o.weight));
        heap.addAll(graph.edges);//M条边，O(logM)
        while (!heap.isEmpty()) {//M条边
            Edge edge = heap.poll();//O(logM)
            if (!u.isSameSet(edge.from, edge.to)) {//O(1)
                u.union(edge.from, edge.to);
                ans.add(edge);
            }
        }
        return ans;
    }
}

